package memory

import (
	"context"
	"time"

	"xagent/internal/messages"

	"github.com/google/uuid"
)

// MemoryType represents different types of memory
type MemoryType string

const (
	// Memory types
	MemoryTypeConversation MemoryType = "conversation"
	MemoryTypeKnowledge    MemoryType = "knowledge"
	MemoryTypeTask         MemoryType = "task"
	MemoryTypeSystem       MemoryType = "system"
	MemoryTypeResult       MemoryType = "result"
	MemoryTypeCode         MemoryType = "code"
	MemoryTypeFact         MemoryType = "fact"
	MemoryTypePrompt       MemoryType = "prompt"
	MemoryTypeShortTerm    MemoryType = "short_term"
	MemoryTypeLongTerm     MemoryType = "long_term"
)

// MemoryEntry represents a memory entry
type MemoryEntry struct {
	ID        string                 `json:"id"`
	Type      MemoryType             `json:"type"`
	Content   string                 `json:"content"`
	Metadata  map[string]interface{} `json:"metadata"`
	Embedding []float32              `json:"embedding,omitempty"`
	CreatedAt time.Time              `json:"created_at"`
	UpdatedAt time.Time              `json:"updated_at"`
	Timestamp time.Time              `json:"timestamp"`
	Tags      []string               `json:"tags"`
}

// MemoryStore defines the interface for memory storage
type MemoryStore interface {
	Store(ctx context.Context, memory *MemoryEntry) error
	Get(ctx context.Context, id string) (*MemoryEntry, error)
	Search(ctx context.Context, query string, limit int) ([]*MemoryEntry, error)
	ListByType(ctx context.Context, memoryType MemoryType, limit int) ([]*MemoryEntry, error)
	Delete(ctx context.Context, id string) error
	Clear(ctx context.Context) error
}

// Memory represents an agent's memory interface
type Memory interface {
	// SaveMessage saves a message to memory
	SaveMessage(ctx context.Context, msg messages.Message) error

	// GetMessages retrieves messages from memory
	GetMessages(ctx context.Context) ([]messages.Message, error)

	// GetMetadata returns memory metadata
	GetMetadata() map[string]interface{}

	// GetTags returns memory tags
	GetTags() []string

	// GetType returns the type of memory
	GetType() MemoryType

	// GetContent returns the content of memory
	GetContent() string
}

// MemoryManager manages agent memories
type MemoryManager interface {
	MemoryStore
	// GetMemoryByID returns memory for an agent
	GetMemoryByID(id string) Memory

	// AddToShortTerm adds memory to short term storage
	AddToShortTerm(ctx context.Context, memory *MemoryEntry) error

	// GetByTags retrieves memories by tags
	GetByTags(ctx context.Context, tags []string) ([]*MemoryEntry, error)

	// GetByTimeRange retrieves memories within a time range
	GetByTimeRange(ctx context.Context, start, end time.Time) ([]*MemoryEntry, error)

	// UpdateTags updates memory tags
	UpdateTags(ctx context.Context, id string, tags []string) error

	// GetStats retrieves memory statistics
	GetStats() (*MemoryStats, error)

	// GetRelevantMemories retrieves relevant memories based on query
	GetRelevantMemories(ctx context.Context, query string, limit int, minSimilarity float32) ([]*MemoryEntry, error)

	// Initialize initializes the memory manager
	Initialize(ctx context.Context) error

	// Close closes the memory manager
	Close() error
}

// MemoryStats represents memory usage statistics
type MemoryStats struct {
	TotalCount     int64
	TypeCounts     map[MemoryType]int64
	AverageLatency time.Duration
	UsedSpace      int64
}

// Specialized memory types

// ConversationMemory represents a conversation history
type ConversationMemory struct {
	MemoryEntry           // 使用 MemoryEntry 而不是 Memory 接口
	Messages    []Message // conversation messages
}

// Message represents a message in a conversation
type Message struct {
	Role      string    // role of the sender (e.g., "user", "assistant")
	Content   string    // message content
	Timestamp time.Time // when the message was sent
}

// CodeMemory represents stored code information
type CodeMemory struct {
	MemoryEntry        // 使用 MemoryEntry
	Language    string // programming language
	Code        string // actual code
	Context     string // context or description
}

// FactMemory represents a stored fact or piece of knowledge
type FactMemory struct {
	MemoryEntry            // 使用 MemoryEntry
	Fact        string     // the actual fact
	Tags        []string   // related tags
	Confidence  float32    // confidence score
	Source      string     // where this fact came from
	ValidUntil  *time.Time // when this fact expires (if applicable)
}

// PromptMemory represents a stored prompt template
type PromptMemory struct {
	MemoryEntry                     // 使用 MemoryEntry
	Template    string              // prompt template with placeholders
	Variables   []string            // list of variables in the template
	Examples    []map[string]string // example variable values
	Usage       int                 // number of times this prompt was used
	Performance map[string]float32  // performance metrics
}

// ResultMemory represents the result of an operation
type ResultMemory struct {
	MemoryEntry                        // 使用 MemoryEntry
	Operation   string                 // operation that produced this result
	Input       map[string]interface{} // input parameters
	Output      interface{}            // operation output
	Duration    time.Duration          // how long the operation took
	Success     bool                   // whether the operation succeeded
	Error       string                 // error message if operation failed
}

// SystemMemory represents system-level information
type SystemMemory struct {
	MemoryEntry                       // 使用 MemoryEntry
	Category    string                // system category (e.g., "config", "state", "metric")
	Key         string                // unique key within category
	Value       interface{}           // stored value
	Persistent  bool                  // whether this should persist across restarts
	LastUpdated time.Time             // when the value was last updated
	History     []SystemMemoryHistory // value history if tracking is enabled
}

// SystemMemoryHistory represents a historical value in system memory
type SystemMemoryHistory struct {
	Value     interface{} // historical value
	Timestamp time.Time   // when this value was recorded
	Reason    string      // why the value changed
}

// Helper functions for type conversion
func AsConversationMemory(m *MemoryEntry) (*ConversationMemory, bool) {
	if m.Type != MemoryTypeConversation {
		return nil, false
	}
	conv := &ConversationMemory{MemoryEntry: *m}
	if msgs, ok := m.Metadata["messages"].([]Message); ok {
		conv.Messages = msgs
	}
	return conv, true
}

func AsCodeMemory(m *MemoryEntry) (*CodeMemory, bool) {
	if m.Type != MemoryTypeCode {
		return nil, false
	}
	code := &CodeMemory{MemoryEntry: *m}
	if lang, ok := m.Metadata["language"].(string); ok {
		code.Language = lang
	}
	if content, ok := m.Metadata["code"].(string); ok {
		code.Code = content
	}
	if ctx, ok := m.Metadata["context"].(string); ok {
		code.Context = ctx
	}
	return code, true
}

func AsFactMemory(m *MemoryEntry) (*FactMemory, bool) {
	if m.Type != MemoryTypeFact {
		return nil, false
	}
	fact := &FactMemory{MemoryEntry: *m}
	if f, ok := m.Metadata["fact"].(string); ok {
		fact.Fact = f
	}
	if tags, ok := m.Metadata["tags"].([]string); ok {
		fact.Tags = tags
	}
	if conf, ok := m.Metadata["confidence"].(float32); ok {
		fact.Confidence = conf
	}
	return fact, true
}

func AsPromptMemory(m *MemoryEntry) (*PromptMemory, bool) {
	if m.Type != MemoryTypePrompt {
		return nil, false
	}
	prompt := &PromptMemory{MemoryEntry: *m}
	if tmpl, ok := m.Metadata["template"].(string); ok {
		prompt.Template = tmpl
	}
	if vars, ok := m.Metadata["variables"].([]string); ok {
		prompt.Variables = vars
	}
	return prompt, true
}

func AsResultMemory(m *MemoryEntry) (*ResultMemory, bool) {
	if m.Type != MemoryTypeResult {
		return nil, false
	}
	result := &ResultMemory{MemoryEntry: *m}
	if op, ok := m.Metadata["operation"].(string); ok {
		result.Operation = op
	}
	if dur, ok := m.Metadata["duration"].(time.Duration); ok {
		result.Duration = dur
	}
	if success, ok := m.Metadata["success"].(bool); ok {
		result.Success = success
	}
	return result, true
}

func AsSystemMemory(m *MemoryEntry) (*SystemMemory, bool) {
	if m.Type != MemoryTypeSystem {
		return nil, false
	}
	sys := &SystemMemory{MemoryEntry: *m}
	if cat, ok := m.Metadata["category"].(string); ok {
		sys.Category = cat
	}
	if key, ok := m.Metadata["key"].(string); ok {
		sys.Key = key
	}
	if val, ok := m.Metadata["value"]; ok {
		sys.Value = val
	}
	return sys, true
}

// MemoryConfig contains configuration for memory stores
type MemoryConfig struct {
	Type    string            // type of memory store (e.g., "redis", "postgres", "vector")
	Address string            // address of the memory store
	Options map[string]string // additional store-specific options
}

// Add these conversion functions
func NewMemory(entry *MemoryEntry) messages.Message {
	return &messages.BasicMessage{
		MessageID:  entry.ID,
		MsgType:    messages.MessageType(entry.Type),
		MsgContent: entry.Content,
		MsgMetadata: map[string]interface{}{
			"role": entry.Metadata["role"],
		},
	}
}

func ToMemoryEntry(msg messages.Message) *MemoryEntry {
	return &MemoryEntry{
		ID:      uuid.New().String(),
		Type:    MemoryType(msg.GetType()),
		Content: msg.GetContent(),
		Metadata: map[string]interface{}{
			"role": msg.GetRole(),
		},
		Timestamp: time.Now(),
		CreatedAt: time.Now(),
		UpdatedAt: time.Now(),
	}
}

// BaseMemory provides a basic implementation of Memory interface
type BaseMemory struct {
	manager  MemoryManager
	id       string
	metadata map[string]interface{}
	tags     []string
	content  string
	memType  MemoryType
}

func NewBaseMemory(manager MemoryManager, id string, metadata map[string]interface{}) *BaseMemory {
	return &BaseMemory{
		manager:  manager,
		id:       id,
		metadata: metadata,
		tags:     make([]string, 0),
		memType:  MemoryTypeSystem, // 默认类型
		content:  "",
	}
}

func (m *BaseMemory) SaveMessage(ctx context.Context, msg messages.Message) error {
	entry := ToMemoryEntry(msg)
	return m.manager.Store(ctx, entry)
}

func (m *BaseMemory) GetMessages(ctx context.Context) ([]messages.Message, error) {
	entries, err := m.manager.ListByType(ctx, MemoryTypeConversation, 100)
	if err != nil {
		return nil, err
	}
	msgs := make([]messages.Message, len(entries))
	for i, entry := range entries {
		msgs[i] = NewMemory(entry)
	}
	return msgs, nil
}

func (m *BaseMemory) GetMetadata() map[string]interface{} {
	return m.metadata
}

func (m *BaseMemory) GetTags() []string {
	return m.tags
}

// 添加设置标签的方法
func (m *BaseMemory) SetTags(tags []string) {
	m.tags = tags
}

// HandleMessage processes a message
func (b *BaseMemoryManager) HandleMessage(msg messages.Message) error {
	// Process message content
	entry := &MemoryEntry{
		ID:      uuid.NewString(),
		Type:    MemoryTypeConversation,
		Content: msg.GetContent(),
		Metadata: map[string]interface{}{
			"role": msg.GetRole(),
		},
		CreatedAt: time.Now(),
		UpdatedAt: time.Now(),
		Timestamp: time.Now(),
	}
	return b.Store(context.Background(), entry)
}

// MemoryEntry 需要实现 Memory 接口
func (m *MemoryEntry) GetType() MemoryType {
	return m.Type
}

func (m *MemoryEntry) GetContent() string {
	return m.Content
}

func (m *MemoryEntry) GetMetadata() map[string]interface{} {
	return m.Metadata
}

func (m *MemoryEntry) GetTags() []string {
	return m.Tags
}

// 实现 GetType 方法
func (m *BaseMemory) GetType() MemoryType {
	return m.memType
}

// 实现 GetContent 方法
func (m *BaseMemory) GetContent() string {
	return m.content
}

// 添加设置方法
func (m *BaseMemory) SetType(memType MemoryType) {
	m.memType = memType
}

func (m *BaseMemory) SetContent(content string) {
	m.content = content
}
